from PIL import Image, ImageDraw, ImageFont
import numpy as np
import os
import random

# 配置
FONTS_DIR = "fonts"  # 字体文件存放目录
TRAIN_OUTPUT_DIR = "data/train"  # 训练集保存目录
VAL_OUTPUT_DIR = "data/val"  # 验证集保存目录
TEXT_POOL = "晨曦微露小镇苏醒街头巷尾人群如织孩童嬉戏笑声飞扬公园绿树成荫面包香气扑鼻老人闲坐树下聊天山峦连绵蓝天如洗生活节奏缓慢惬意平凡中蕴含温暖幸福悄然流淌在每一个角落时光静好岁月无声小镇的故事在每一个清晨和黄昏中继续着每一个瞬间都值得珍惜每一份情感都深沉而绵长"  # 包含100个字符的字符串
NUM_TRAIN_IMAGES_PER_FONT = 1000  # 每种字体生成的训练集图片数量
NUM_VAL_IMAGES_PER_FONT = 100  # 每种字体生成的验证集图片数量
IMAGE_SIZE = (224, 224)  # 图片尺寸
BACKGROUND_COLORS = ["white", "lightgray", "lightblue", "lightyellow"]  # 背景颜色
MAX_NOISE_LEVEL = 0.02  # 最大噪声水平（0-1）

# 确保目录存在
os.makedirs(FONTS_DIR, exist_ok=True)
os.makedirs(TRAIN_OUTPUT_DIR, exist_ok=True)
os.makedirs(VAL_OUTPUT_DIR, exist_ok=True)

# 加载字体文件
font_files = [f for f in os.listdir(FONTS_DIR) if f.endswith(".ttf")]
if not font_files:
    raise FileNotFoundError(f"请在 {FONTS_DIR} 目录下放置字体文件（如 .ttf）")

# 添加局部噪声
def add_local_noise(image, noise_level):
    img_array = np.array(image)
    mask = (img_array == np.array([255, 255, 255])).all(axis=-1)  # 找到背景区域
    noise = np.random.normal(0, noise_level * 255, img_array.shape).astype(np.uint8)
    img_array[mask] = np.clip(img_array[mask] + noise[mask], 0, 255)
    return Image.fromarray(img_array)

# 绘制多行文字
def draw_multiline_text(draw, text, font, max_width, start_x, start_y, text_color, line_spacing=10):
    lines = []
    words = text.split()
    current_line = words[0]
    for word in words[1:]:
        test_line = current_line + " " + word
        if draw.textlength(test_line, font=font) <= max_width:
            current_line = test_line
        else:
            lines.append(current_line)
            current_line = word
    lines.append(current_line)
    
    y = start_y
    for line in lines:
        draw.text((start_x, y), line, font=font, fill=text_color)
        y += font.size + line_spacing

# 生成图片
def generate_images(font_path, output_dir, num_images):
    font_name = os.path.splitext(os.path.basename(font_path))[0]  # 字体名称（如“宋体”）
    font_output_dir = os.path.join(output_dir, font_name)
    os.makedirs(font_output_dir, exist_ok=True)
    
    for i in range(num_images):
        # 创建背景
        bg_color = random.choice(BACKGROUND_COLORS)
        img = Image.new("RGB", IMAGE_SIZE, color=bg_color)
        draw = ImageDraw.Draw(img)
        
        # 随机化字体大小
        font_size = random.randint(30, 50)
        font = ImageFont.truetype(font_path, size=font_size)
        
        # 从字符串中随机选取1-5个字符
        num_chars = random.randint(1, 5)  # 随机选取1到5个字符
        text = ''.join(random.sample(TEXT_POOL, num_chars))  # 从TEXT_POOL中随机选取字符
        
        # 随机化文字位置
        text_x = random.randint(10, IMAGE_SIZE[0] - 100)
        text_y = random.randint(10, IMAGE_SIZE[1] - 50)
        
        # 随机化文字颜色
        text_color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
        
        # 绘制多行文字
        draw_multiline_text(draw, text, font, max_width=IMAGE_SIZE[0] - 20, start_x=text_x, start_y=text_y, text_color=text_color)
        
        # 添加局部噪声
        noise_level = random.uniform(0, MAX_NOISE_LEVEL)
        if noise_level > 0:
            img = add_local_noise(img, noise_level)
        
        # 保存图片
        img.save(os.path.join(font_output_dir, f"{i}.png"))

# 生成训练集和验证集
for font_file in font_files:
    font_path = os.path.join(FONTS_DIR, font_file)
    
    # 生成训练集
    generate_images(font_path, TRAIN_OUTPUT_DIR, NUM_TRAIN_IMAGES_PER_FONT)
    
    # 生成验证集
    generate_images(font_path, VAL_OUTPUT_DIR, NUM_VAL_IMAGES_PER_FONT)

print(f"训练集生成完成，保存到 {TRAIN_OUTPUT_DIR}")
print(f"验证集生成完成，保存到 {VAL_OUTPUT_DIR}")